'''
Created on May 3, 2012

@author: analyst
'''
from teSdk import *
import os
import pefile
import datetime
import re

g_DEUBG = 1

class deUnpackZA():
    
    def __init__(self):
        
        # Function prototypes like C for breakpoint callback functions
        self.CbOnEntryPoint     = fBreakPoint(self.OnEntryPoint)
        self.CbOnVirtualProtect = fBreakPoint(self.OnVirtualProtect)
        self.CbOnOepCode        = fBreakPoint(self.OnOepCode)
        self.CbOnExecuteOep     = fBreakPoint(self.OnExecuteOep)
        self.CbOnWriteImgBase   = fBreakPoint(self.OnWriteImgBase)
        self.CbOnWriteImgBaseEP = fBreakPoint(self.OnWriteImgBaseEP)
        
        # Win32 API function definition
        self.krnl32BaseAddr = windll.kernel32.GetModuleHandleA("kernel32.dll")
        self.VirtualProtect = windll.kernel32.GetProcAddress(self.krnl32BaseAddr, "VirtualProtect")
        self.ReadProcessMemory = windll.kernel32.ReadProcessMemory
        
        self.wild = c_ubyte(0xcc)

        self.lo   = 0
        self.hi   = 0
        
        self.ep1    = c_ulong(0)
        self.ep2    = c_ulong(0)
        
        self.PAGE_EXECUTE_READ_WRITE = 0x40
        
        '''
        CALL EAX    FF D0
        CALL EBX    FF D3
        CALL ECX    FF D1
        CALL EDX    FF D2
        CALL EDI    FF D7
        CALL ESI    FF D6
        CALL ESP    FF D4
        CALL ESP    FF D5
        '''
        self.inst = [(c_ubyte * 2)(0xFF, 0xD0), 
                     (c_ubyte * 2)(0xFF, 0xD1), 
                     (c_ubyte * 2)(0xFF, 0xD2), 
                     (c_ubyte * 2)(0xFF, 0xD3)]
        
        '''
        Variables that hold the result of the unpacking process
        '''
        self.totalfile = 0
        self.fileprocessed = 0
        self.fileunprocessed = 0
        self.fileunpacked = 0
    
    def OnExecuteOep(self):
        oep         = TE.GetContextData(UE_EIP)
        hProcess    = self.ProcInfo[0].hProcess
   
    
        print("   [+] Found OEP!")
        self.logfile.write("   [+] Found OEP!")
        print("   [+] OEP address: 0x%08x.\n" % oep)
        self.logfile.write("   [+] OEP address: 0x%08x.\n" % oep)
        
        if TE.DumpProcess(hProcess, self.Base, self.Output, oep):
            print("   [+] Process dumped as: %s " % self.Output)
            self.logfile.write("   [+] Process dumped as: %s\n" % self.Output)
            print("[+] Done!")
            self.logfile.write("[+] Done!\n")
            self.fileunpacked += 1
        else:
            print("   [-] Failed to dump process.")
            self.logfile.write("   [-] Failed to dump process.\n")
        
        TE.StopDebug()
        
    def OnWriteImgBaseEP(self):
        hProcess        = self.ProcInfo[0].hProcess
        bytesRead       = c_void_p()
        base            = self.Base
        offset_to_oep   = self.offset_to_ep
        
        print("   [+] Hardware write breakpoint hit on 0x%08x" % (base + offset_to_oep))
        self.logfile.write("   [+] Hardware write breakpoint hit on 0x%08x.\n" % (base + offset_to_oep))
        self.ReadProcessMemory(hProcess, base + offset_to_oep, byref(self.ep2), 4, byref(bytesRead))
        print("      [+] Entry point (2): 0x%08x" % self.ep2.value)
        self.logfile.write("      [+] Entry point (2): 0x%08x.\n" % self.ep2.value)
        if self.ep2.value != 0 and self.ep1.value == self.ep2.value:
            TE.DeleteHardwareBreakPoint(self.Dr)
            oep = self.ep1.value
            if TE.SetHardwareBreakPoint(base + oep, self.Dr, UE_HARDWARE_EXECUTE, UE_HARDWARE_SIZE_1, self.CbOnExecuteOep):
                print("      [+] Hardware breakpoint set on 0x%08x" % (base + oep))
                self.logfile.write("      [+] Hardware breakpoint set on 0x%08x.\n" % (base + oep))
            else:
                print("      [-] Cannot set hardware breakpoint")
                self.logfile.write("      [-] Cannot set hardware breakpoint.\n")
                TE.StopDebug()
                
    def OnWriteImgBase(self):
        esi         = TE.GetContextData(UE_ESI)
        hProcess    = self.ProcInfo[0].hProcess
        pe          = c_ulong(0)
        bytesRead   = c_void_p()
        
        TE.DeleteHardwareBreakPoint(self.Dr)
        print("   [+] Hardware write breakpoint hit on 0x%08x" % self.Base)
        self.logfile.write("   [+] Hardware write breakpoint hit on 0x%08x.\n" % self.Base)
        print("      [+] Base address of unpacked code: 0x%08x" % esi)
        self.logfile.write("      [+] Base address of unpacked code: 0x%08x.\n" % esi)
        self.ReadProcessMemory(hProcess, esi + 0x3C, byref(pe), 4, byref(bytesRead))
        self.offset_to_ep = pe.value + 0x28
        self.ReadProcessMemory(hProcess, esi + self.offset_to_ep, byref(self.ep1), 4, byref(bytesRead))
        
        print("      [+] Entry point (1): 0x%08x" % self.ep1.value)
        self.logfile.write("      [+] Entry point (1): 0x%08x.\n" % self.ep1.value)
        if self.ep1.value != 0:
                base = self.Base
                if TE.SetHardwareBreakPoint(base + self.offset_to_ep, self.Dr, UE_HARDWARE_WRITE, UE_HARDWARE_SIZE_4, self.CbOnWriteImgBaseEP):
                    print("      [+] Hardware breakpoint set on 0x%08x" % (base + self.offset_to_ep))
                    self.logfile.write("      [+] Hardware breakpoint set on 0x%08x.\n" % (base + self.offset_to_ep))
                else:
                    print("      [-] Cannot set hardware breakpoint")
                    self.logfile.write("      [-] Cannot set hardware breakpoint.\n")
                    TE.StopDebug()
            
        
    def OnOepCode(self):
        
        hProcess    = self.ProcInfo[0].hProcess
        pe          = c_ulong()
        ep          = c_ulong()
        bytesRead   = c_void_p()
        eip         = TE.GetContextData(UE_EIP)
        ins = TE.Disassemble(eip)
        print("   [+] Breakpoint on instruction that jump to OEP hit: 0x%08x" % eip)
        self.logfile.write("   [+] Breakpoint on instruction that jump to OEP hit: 0x%08x.\n" % eip)
        if ins.find("CALL EAX") != -1: oep = TE.GetContextData(UE_EAX)
        elif ins.find("CALL EBX") != -1: oep = TE.GetContextData(UE_EBX)
        elif ins.find("CALL ECX") != -1: oep = TE.GetContextData(UE_ECX)
        elif ins.find("CALL EDX") != -1: oep = TE.GetContextData(UE_EDX)
        
        self.ReadProcessMemory(hProcess, self.Base + 0x3C, byref(pe), 4, byref(bytesRead))
        self.ReadProcessMemory(hProcess, self.Base + pe.value + 0x28, byref(ep), 4, byref(bytesRead))     
       
        if oep == (self.Base + ep.value):
            print("   [+] Found OEP!")
            self.logfile.write("   [+] Found OEP!\n")
            print("   [+] OEP address: 0x%08x" % oep)
            self.logfile.write("   [+] OEP address: 0x%08x.\n" % oep)
            if TE.DumpProcess(self.ProcInfo[0].hProcess, self.Base, self.Output, oep):
                print("   [+] Process dumped as: %s " % self.Output)
                self.logfile.write("   [+] Process dumped as: %s \n" % self.Output)
                print("[+] Done!")
                self.logfile.write("[+] Done!\n")
                self.fileunpacked += 1
            else:
                print("   [-] Failed to dump process.")
                self.logfile.write("   [-] Failed to dump process.\n")
    
        TE.StopDebug()
        
        
    def OnVirtualProtect(self):
        esp             = TE.GetContextData(UE_ESP)
        RetnAddress     = esp
        BaseAddress     = esp + 4
        dwSizeAddress   = esp + 8
        flNewProtectAddr = esp + 12
        targetaddr      = c_void_p()
        retnaddr        = c_void_p()
        dwSize          = c_void_p()
        flNewProtect    = c_void_p()
        buffsz          = 4
        bytesRead       = c_ulong(0)
        hProcess        = self.ProcInfo[0].hProcess

        self.ReadProcessMemory(hProcess, RetnAddress, byref(retnaddr), buffsz, byref(bytesRead))
        self.ReadProcessMemory(hProcess, BaseAddress, byref(targetaddr), buffsz, byref(bytesRead))
        self.ReadProcessMemory(hProcess, dwSizeAddress, byref(dwSize), buffsz, byref(bytesRead))
        self.ReadProcessMemory(hProcess, flNewProtectAddr, byref(flNewProtect), buffsz, byref(bytesRead))
        
        if dwSize.value == self.SizeOfImg and flNewProtect.value == self.PAGE_EXECUTE_READ_WRITE:
            '''
            TYPE 1
            example SHA1 - 27422131689a3076118f6dfb3d037a2d3e07ae7f
            example SHA1 - 0d1b141351a909d48f39946894d9c673e71ec57c
            
            * Check BaseAddress which should be equivalent to the original executable image base
            * Check dwSize
            * Check flNewProtect which should be PAGE_EXECUTE_READWRITE
            '''
            print("   [+] Memory address where protection will be updated: 0x%x" % targetaddr.value)
            self.logfile.write("   [+] Memory address where protection will be updated: 0x%x.\n" % targetaddr.value)
            print("      [+] Memory size to be updated: 0x%08x" % dwSize.value)
            self.logfile.write("      [+] Memory size to be updated: 0x%08x.\n" % dwSize.value)
            print("      [+] New memory protection: %s" % "PAGE_EXECUTE_READ_WRITE")
            self.logfile.write("      [+] New memory protection: %s.\n" % "PAGE_EXECUTE_READ_WRITE")
            
            ###
            ### The retn address should not fall inside the loaded image address
            ###
            if not (retnaddr.value >= self.Base and retnaddr.value <= (self.Base + self.SizeOfImg)):
                print("      [+] Address of the unpacking routine: 0x%08x" % retnaddr.value)
                self.logfile.write("      [+] Address of the unpacking routine: 0x%08x.\n" % retnaddr.value)
                base = retnaddr.value
                for i in range(len(self.inst)):
                    offset = 0
                    # TODO: Fix the search block size. Make it dynamic
                    self.loc = TE.Find(base, 0x500, byref(self.inst[i]), 2, 0)
                    while self.loc != 0:
                        if TE.SetBPX(self.loc, UE_BREAKPOINT, self.CbOnOepCode):
                            print("      [+] Breakpoint set on 0x%08x" % self.loc)
                            self.logfile.write("      [+] Breakpoint set on 0x%08x.\n" % self.loc)
                        else:
                            print("      [-] Cannot set breakpoint on 0x%08x" % self.loc)
                            self.logfile.write("      [-] Cannot set breakpoint on 0x%08x.\n" % self.loc)
                            TE.StopDebug()
                        offset = self.loc - base
                        self.loc = TE.Find(base + offset, 0x500, byref(self.inst[i]), 2, 0)
                        
            '''
            TYPE 2
            example SHA1 - ba75a60aeeaa151dcf2d4c44091bd35a4a06e75e
            example SHA1 - ae42f6f5e8ca769a299a45e83cc991e7d0b219b7
            
            * Call LocalAlloc(LMEM_ZEROINIT, heap_size) => The allocated heap should contain the unpacking routine
            
            * only call 2 VirtualProtect win32 API
                o Changed the memory protection of the loaded image itself with random size? (not SizeOfImage)
                o Changed the memory protection of the heap that has the unpacking routine which contains a bunch of junk codes
                   + Set HW BP on the Image base of where the memory protection will be changed
                   + When the HW BP is hit, noticed that the ESI register contains the unpacked executable
                    ~ Parse the PE header here and get the __EP1__
                    ~ At the same time, get the __offset__ to the EP 
                    ~ Set the HW write access breakpoint at IMAGE_BASE + __OFFSET__ and read the __EP2__ from there
                    ~ Compare both __EP1__ and __EP2__
                       # If there are the same, set hardware execution breakpoint at IMAGE_BASE + __EP__
                       # OEP will be hit
            '''
        elif targetaddr.value == self.Base and flNewProtect.value == self.PAGE_EXECUTE_READ_WRITE:
            base = c_void_p(targetaddr.value)
            self.Dr = c_ulong(0)
            TE.GetUnusedHardwareBreakPointRegister(byref(self.Dr))
            if TE.SetHardwareBreakPoint(base.value, self.Dr, UE_HARDWARE_WRITE, UE_HARDWARE_SIZE_4, self.CbOnWriteImgBase):
                print("   [+] Hardware breakpoint set on 0x%08x" % base.value)
                self.logfile.write("   [+] Hardware breakpoint set on 0x%08x.\n" % base.value)
            else:
                print("   [-] Cannot set hardware breakpoint")
                self.logfile.write("   [-] Cannot set hardware breakpoint.\n")
                TE.StopDebug()
                
    def OnEntryPoint(self):
        
        self.Base   = TE.GetDebuggedFileBaseAddress()
        if g_DEUBG: print("   [+] EIP: 0x%08x" % TE.GetContextData(UE_EIP))
        if g_DEUBG: print("   [+] Image loaded in: 0x%08x" %self.Base)
        if g_DEUBG: print("   [+] EntryPoint: 0x%08x" % self.EP)
        if g_DEUBG: print("   [+] hThread: 0x%08x" % self.ProcInfo[0].hThread)
        if g_DEUBG: print("   [+] hProcess: 0x%08x" % self.ProcInfo[0].hProcess)
        if g_DEUBG: print("   [+] VirtualProtect address: 0x%08x" % self.VirtualProtect)
        
        ###
        ### Set breakpoint on VirtualProtect
        ###
        if TE.SetBPX(self.VirtualProtect, UE_BREAKPOINT, self.CbOnVirtualProtect):
            print("   [+] Breakpoint set on VirtualProtect")
            self.logfile.write("   [+] Breakpoint set on VirtualProtect.\n")
        else:
            print("   [-] Failed to set breakpoint on VirtualProtect")
            self.logfile.write("   [-] Failed to set breakpoint on VirtualProtect.\n")
            TE.StopDebug()
    
    def GetNumberOfImportedDll(self):
        count = 0
        try:
            for dllentry in self.pe.DIRECTORY_ENTRY_IMPORT:
                count += 1
                if dllentry.dll == "mscoree.dll":
                    return -1
        except:
            print("[-] Error parsing import table")
            self.logfile.write("[-] Error parsing import table.\n")
            exit()
            
        return count
        
    def OnInitialize(self, pOriginalFn):
        
        self.Input  = pOriginalFn
        # TODO: Files might not have extension name, eg: full_path_name\n32 => full_path_name\unpacked\n32
        self.Output = pOriginalFn[:len(pOriginalFn) - 3] + "unpacked" + pOriginalFn[len(pOriginalFn) - 4:]
        
        self.Validity = FILE_STATUS_INFO()
        if not TE.IsPE32FileValidEx(pOriginalFn, UE_DEPTH_DEEP, byref(self.Validity)) \
           or  self.Validity.OveralEvaluation != UE_RESULT_FILE_OK:
            raise Exception
            return
        
        # Before initializing the debugger, we will check if the unpacker support the EXE
        if self.GetNumberOfImportedDll() < 0:
            print("[-] .NET packer is not supported")
            self.logfile.write("[-] .NET packer is not supported.\n")
            
        # InitDebugEx is a wrapper code that will execute CreateProcess Win32 API
        self.ProcInfo = TE.InitDebugEx(pOriginalFn, 0, 0, self.CbOnEntryPoint)
        
        if self.ProcInfo:
            print("[+] Debugger initialized.")
            self.logfile.write("[+] Debugger initialized.\n")

            self.GotTLS     = TE.GetPE32Data(pOriginalFn, 0, UE_TLSTABLEADDRESS)
            self.TextSz     = TE.GetPE32Data(pOriginalFn, 0, UE_SECTIONVIRTUALSIZE)
            self.TextVA     = TE.GetPE32Data(pOriginalFn, 0, UE_SECTIONVIRTUALOFFSET)
            self.SizeOfImg  = TE.GetPE32Data(pOriginalFn, 0, UE_SIZEOFIMAGE)
            self.ImgBase    = TE.GetPE32Data(pOriginalFn, 0, UE_IMAGEBASE)
            self.EP         = TE.GetPE32Data(pOriginalFn, 0, UE_OEP)

            self.SnapRange  = TE.GetPE32Data(pOriginalFn, 0, UE_SECTIONVIRTUALOFFSET)
            self.SnapSize   = self.EP - self.SnapRange

            TE.DebugLoop() # This function will trigger the OnEntryPoint callback function
        else:
            print("[-] Could not initialize debugger")
            self.logfile.write("[-] Could not initialize debugger.\n")
        
    def StartUnpack(self):
        if __name__ == '__main__':
            
            today = datetime.datetime.today()
            timestamp = "{0}{1}{2}{3}{4}{5}"
            try:
                if os.path.split(os.sys.argv[0])[0] == "":
                    logpath = os.curdir + os.sep
                else:
                    logpath = os.path.split(os.sys.argv[0])[0] + os.sep
                self.logfile = open(logpath  + "deUnpackZA" + 
                               timestamp.format(today.year, today.month, today.day, today.hour, today.minute, today.microsecond) + ".log", 
                               "w")
            except:
                print("[-] Could not create log file")
                exit()
                
            path = ""
            file = ""
            if len(os.sys.argv) < 2:
                print("***********************************************************************")
                print("Generic Unpacker for ZeroAccess (GUZA)")
                print("")
                print("NOTE: Please ONLY run this script in a testing environment (eg: VMware)")
                print("***********************************************************************")
                print("")
                print("Usage: %s <path_to_scan>" % os.sys.argv[0])
                exit()
            elif os.path.isdir(os.sys.argv[1]):
                path = os.sys.argv[1]
            elif os.path.isfile(os.sys.argv[1]):
                file = os.sys.argv[1]
            
            if path != "":
                numfile = 1
                for dirpath, dirnames, files in os.walk(path):
                    print("[+] Unpacking files in: " + dirpath + " ...")
                    self.logfile.write("[+] Unpacking files in: " + dirpath + " ...\n")
                    for file in files:
                        try:
                            self.pe = pefile.PE(path+"\\"+file)
                            if self.pe.is_exe():
                                print("[%d] Unpacking \"%s\" ..." % (numfile, file))
                                self.logfile.write("[%d] Unpacking \"%s\" ...\n" % (numfile, file))
                                numfile += 1
                                self.OnInitialize(path+"\\"+file)
                            else:
                                print("[%d] Skipping non EXE file \"%s\"" % (numfile, file))
                                self.logfile.write("[%d] Skipping non EXE file \"%s\"\n" % (numfile, file))
                                numfile += 1
                            self.fileprocessed += 1
                        except:
                            print("[-] Unable to process %s. It is probably not a PE file or corrupted." % file)
                            self.logfile.write("[-] Unable to process %s. It is probably not a PE file or corrupted.\n" % file)
                            self.fileunprocessed += 1
                            continue
            elif file != "":
                    # There is only one file to be unapacked
                    files = [file]
                    try:
                        self.pe = pefile.PE(file)
                        if self.pe.is_exe():
                            print("[+] Unpacking file: " + file + " ...")
                            self.logfile.write("[+] Unpacking file: " + file + " ...\n")
                            self.OnInitialize(file)
                        else:
                            print("[+] Skipping non EXE file " + file)
                            self.logfile.write("[+] Skipping non EXE file " + file + "\n")
                        self.fileprocessed += 1
                    except:
                        print("[-] Unable to read file %s. It is probably not a PE file or corrupted." % file)
                        self.logfile.write("[-] Unable to read file %s. It is probably not a PE file or corrupted.\n" % file)
                        self.fileunprocessed += 1
            else:
                print("Example: %s fullpath\\filename.exe" % os.sys.argv[0])
            
            self.totalfile = len(files)
            print("\nResult:\n")
            self.logfile.write("\nResult:\n")
            print("Total number of files: %d" % self.totalfile)
            self.logfile.write("\nTotal number of files: %d\n" % self.totalfile)
            print("Total number of files processed: %d" % self.fileprocessed)
            self.logfile.write("Total number of files processed: %d\n" % self.fileprocessed)
            print("Total number of files unprocessed: %d" % self.fileunprocessed)
            self.logfile.write("Total number of files unprocessed: %d\n" % self.fileunprocessed)
            print("Total number of unpacked files: %d" % self.fileunpacked)
            self.logfile.write("Total number of unpacked files: %d\n" % self.fileunpacked)
            self.logfile.close()
            
ZAUnpacker = deUnpackZA()
ZAUnpacker.StartUnpack()
    